# !/usr/bin/python3
# -*- coding:utf-8 -*-
# Copyright 2021 The Chinaunicom Software Team. All rights reserved.
# @Author : dyu
# @Time   : 2021-12-22

### 模型推理模块 ###

import os
from src.intelligent_interaction.engine.sem_conf import Configuration
import tensorflow as tf
import pandas as pd
import warnings, time
warnings.simplefilter(action='ignore', category=FutureWarning)

from tensorflow.python.keras.backend import set_session
from src.intelligent_interaction.engine.sem_predictor import Predictor
from src.intelligent_interaction.engine.annoy_model import AnnoyRecall
from src.intelligent_interaction.engine.make_dataset import word_segment
from src.intelligent_interaction.engine.sem_utils import _read_json
config = Configuration()
memoryList = list(map(int, os.popen("nvidia-smi -q -d Memory | grep -A4 GPU | grep Total | awk '{print $3}'").readlines()))
memoryTotal = memoryList[0]
memoryLimited = 5000.0  # use 5G memory


class BasePredictor(object):
    def __init__(self, task_name):
        # log_server.logging('===============Prediction Process Beginning !===============')
        self.task_name = task_name
        self.faq_df = pd.read_csv(config.LoadParams['all_data'], header=0)
        tf_config = tf.ConfigProto()
        tf_config.gpu_options.per_process_gpu_memory_fraction = memoryLimited / memoryTotal
        self.sess = tf.Session(config=tf_config)
        self.graph = tf.get_default_graph()
        set_session(self.sess)
        self.model, self.ar = self.create_model()
        print('模型加载完毕')

    def create_model(self):
        # log_server.logging('===============Creating model.===============')
        predictor = Predictor(model_path=config.LoadParams['model'])
        ar = AnnoyRecall(f_dim=128)
        return predictor, ar

    def predict(self, question):
        # log_server.logging('Calling model: {}'.format(self.grade_subject))
        start_time = time.time()
        predict_dct = dict()  # All data recall dict

        ### 第一步：分词 ###
        query_seg = word_segment(question)
        end_time_1 = time.time()
        try:
            ### 第二步：cdssm预测 ###
            output_vector = self.model.predictor_single(query_seg)
            end_time_2 = time.time()
        except Exception as e:
            return []
        
        try:
            ### 第三步：Annoy召回 ###
            qids, distance = self.ar.vecter_recall(vector=output_vector[0],
                                                            search_k=-1,
                                                            n_sims=3)
            end_time_3 = time.time()
        except Exception as e:
            qids, distance = [], []
            
        answers = [self.faq_df.at[int(qid), 'answer'] for qid in qids]
        predict = list(zip(answers, distance))
        # print("Predict:", predict)
        return predict


if __name__ == '__main__':
    faq = BasePredictor('encyclopedia')
    res = faq.predict('菲利普是有哪些电影比较出名的')
    print(res)
